
# import
from langchain_chroma import Chroma
from langchain_community.document_loaders import TextLoader
from langchain_community.embeddings.sentence_transformer import (
    SentenceTransformerEmbeddings,
)
from langchain_text_splitters import CharacterTextSplitter
from langchain.embeddings import HuggingFaceBgeEmbeddings

# load the document and split it into chunks
loader = TextLoader("data/paul_graham_essay.txt", encoding='utf-8')
documents = loader.load()

# split it into chunks
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
docs = text_splitter.split_documents(documents)

# create the open-source embedding function
# embedding model: m3e-base
model_name = "moka-ai/m3e-base"
model_kwargs = {'device': 'cuda'}
encode_kwargs = {'normalize_embeddings': True}
embedding = HuggingFaceBgeEmbeddings(
    model_name=model_name,
    model_kwargs=model_kwargs,
    encode_kwargs=encode_kwargs
)
# embedding_function = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")

try:
# load it into Chroma
    db = Chroma.from_documents(docs, embedding)
    print('db loading success')
except Exception as e:
    print('db loadding failued')
    print(e)
# query it
query = "What did the president say about Ketanji Brown Jackson"
docs = db.similarity_search(query)

# print results
print(docs[0].page_content)
print('finish')

# langchain                                0.2.11
# langchain-chroma                         0.1.2
# langchain-community                      0.2.10
# langchain-core                           0.2.23
# langchain-text-splitters                 0.2.2

# chroma-hnswlib                           0.7.6
# chromadb                                 0.5.5